#include <iostream>
#include <winsock2.h>
#include <ws2tcpip.h>
#include <fstream>
#include <string>
#include <windows.h>  

#pragma comment(lib, "ws2_32.lib")
#pragma warning( disable : 4996 ) // inet_addr 方法已经被遗弃了，直接调用会报错，需要禁用该报错

using namespace std;

//#define BUF_MAX_SIZE (1024*1024)
#define BUF_MAX_SIZE 2000
#define SEVER_PORT 8554

// 创建TCP协议的socket，并返回该socket的文件描述符
static int createTcpSocket()
{
    int sockfd;
    int on = 1;
    // 创建IPV4地址的TCP协议的套接字
    sockfd = socket(AF_INET, SOCK_STREAM, 0);
    if (sockfd < 0)
    {
        return -1;
    }
    // 修改套接字选项，允许地址重用，这样在服务器重启的时候可以避免该端口因为被上次关闭还在占用而无法启动
    setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, (const char*)&on, sizeof(on));
    return sockfd;
}

// 创建UDP协议的socket，并返回该socket的文件描述符
static int createUdpSocket()
{
    int sockfd;
    int on = 1;
    // 创建IPV4地址的UDP协议的套接字
    sockfd = socket(AF_INET, SOCK_DGRAM, 0);
    if (sockfd < 0)
    {
        return -1;
    }
    // 修改套接字选项，允许地址重用，这样在服务器重启的时候可以避免该端口因为被上次关闭还在占用而无法启动
    setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, (const char*)&on, sizeof(on));
    return sockfd;
}

int main()
{
    // 存储windows的socket的相关信息
    WSADATA wsaData;
    // 初始化windows的socket的dll
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0)
    {
        cout << "PC Server Socket Start Up Error" << endl;
    }
    // 创建套接字
    int clientSockfd = createTcpSocket();
    if (clientSockfd < 0)
    {
        cout << "faild to create tcp socket" << endl;
        // 清理windows的socket的dll，释放相关资源
        WSACleanup();
        return -1;
    }
    {
        // 创建地址对象
        struct sockaddr_in addr;
        addr.sin_family = AF_INET;
        addr.sin_port = htons(8554);
        addr.sin_addr.s_addr = inet_addr("127.0.0.1");
        // 连接服务端
        if (connect(clientSockfd, (sockaddr*)&addr, sizeof(addr)) == SOCKET_ERROR)
        {
            cout << "connect to addr failed" << endl;
            // 关闭套接字
            closesocket(clientSockfd);
            // 清理windows的socket的dll，释放相关资源
            WSACleanup();
            return -1;
        }
    }

    // 创建套接字
    int rtpSockfd = createUdpSocket(), rtcpSockfd = createUdpSocket();
    if (rtpSockfd < 0 || rtcpSockfd < 0)
    {
        cout << "faild to create udp socket" << endl;
        // 关闭套接字
        closesocket(clientSockfd);
        closesocket(rtpSockfd);
        closesocket(rtcpSockfd);
        // 清理windows的socket的dll，释放相关资源
        WSACleanup();
        return -1;
    }

    {
        // 创建地址对象
        struct sockaddr_in addr;
        addr.sin_family = AF_INET;
        addr.sin_port = htons(10000);
        addr.sin_addr.s_addr = inet_addr("127.0.0.1");
        // 连接服务端
        if (bind(rtpSockfd, (sockaddr*)&addr, sizeof(addr)) == SOCKET_ERROR)
        {
            cout << "connect to addr failed" << endl;
            // 关闭套接字
            closesocket(clientSockfd);
            closesocket(rtpSockfd);
            closesocket(rtcpSockfd);
            // 清理windows的socket的dll，释放相关资源
            WSACleanup();
            return -1;
        }
    }

    {
        // 创建地址对象
        struct sockaddr_in addr;
        addr.sin_family = AF_INET;
        addr.sin_port = htons(10001);
        addr.sin_addr.s_addr = inet_addr("127.0.0.1");
        // 连接服务端
        if (bind(rtcpSockfd, (sockaddr*)&addr, sizeof(addr)) == SOCKET_ERROR)
        {
            cout << "connect to addr failed" << endl;
            // 关闭套接字
            closesocket(clientSockfd);
            closesocket(rtpSockfd);
            closesocket(rtcpSockfd);
            // 清理windows的socket的dll，释放相关资源
            WSACleanup();
            return -1;
        }
    }

    string msg = "OPTIONS rtsp://127.0.0.1:8554 RTSP/1.0\r\n"
        "CSeq: 1\r\n"
        "User - Agent : Lavf58.29.100\r\n";
    send(clientSockfd, msg.c_str(), msg.length(), 0);

    char* recvBuf = (char*)malloc(BUF_MAX_SIZE);
    int recvLen = 0;

    recvLen = recv(clientSockfd, recvBuf, BUF_MAX_SIZE, 0);
    recvBuf[recvLen] = '\0';
    cout << recvBuf << endl;

    msg = "DESCRIBE rtsp://127.0.0.1:8554 RTSP/1.0\r\n"
        "Accept: application / sdp\r\n"
        "CSeq : 2\r\n"
        "User - Agent : Lavf58.29.100\r\n";
    send(clientSockfd, msg.c_str(), msg.length(), 0);

    recvLen = recv(clientSockfd, recvBuf, BUF_MAX_SIZE, 0);
    recvBuf[recvLen] = '\0';

    msg = "SETUP rtsp://127.0.0.1:8554/track0 RTSP/1.0\r\n"
        "Transport: RTP/AVP/UDP;unicast;client_port=10000-10001\r\n"
        "CSeq: 3\r\n"
        "User - Agent : Lavf58.29.100\r\n";
    send(clientSockfd, msg.c_str(), msg.length(), 0);

    recvLen = recv(clientSockfd, recvBuf, BUF_MAX_SIZE, 0);
    recvBuf[recvLen] = '\0';

    msg = "PLAY rtsp://127.0.0.1:8554 RTSP/1.0\r\n"
        "Range: npt = 0.000 -\r\n"
        "CSeq : 5\r\n"
        "User - Agent : Lavf58.29.100\r\n"
        "Session : 66334873\r\n";
    send(clientSockfd, msg.c_str(), msg.length(), 0);

    recvLen = recv(clientSockfd, recvBuf, BUF_MAX_SIZE, 0);
    recvBuf[recvLen] = '\0';

    // 打开输出文件
    ofstream outFile("output.h264", ios::binary);
    if (!outFile.is_open()) {
        cout << "Error: Unable to open output file." << endl;
        return -1; 
    }
    
    // H.264 startcode
    const char startcode[] = {0x00, 0x00, 0x00, 0x01};
    
    while (true) {
        struct sockaddr_in addr;
        int addr_len = sizeof(addr);
        recvLen = recvfrom(rtpSockfd, recvBuf, BUF_MAX_SIZE, 0, 
                         (struct sockaddr*)&addr, &addr_len);
        if (recvLen < 0) {
            cout << "recv error" << endl;
            break;
        }
        if (recvLen <= 12) {
            continue;
        }

        // 解析RTP头
        unsigned char* rtpHeader = (unsigned char*)recvBuf;
        int payloadType = rtpHeader[1] & 0x7F;
        int marker = rtpHeader[1] >> 7;
        int payloadOffset = 12; // 基本RTP头长度
        
        // 处理可能的CSRC扩展
        if (rtpHeader[0] & 0x10) { // 如果有扩展头
            int extensionLength = ntohs(*(uint16_t*)(recvBuf + payloadOffset + 2));
            payloadOffset += 4 + extensionLength * 4;
        }

        // 只处理H.264负载(Payload Type 96)
        if (payloadType == 96) {
            unsigned char* payload = (unsigned char*)(recvBuf + payloadOffset);
            //0x1F : 0001 1111
            int nalType = payload[0] & 0x1F;
            
            // 处理单个NAL单元
            if (nalType <= 23) {
                outFile.write(startcode, sizeof(startcode));
                outFile.write((char*)payload, recvLen - payloadOffset);
            } 
            // 处理分片单元(FU-A)
            else if (nalType == 28) {
                // 0x80: 1000 0000
                // 0x40: 0100 0000
                int startBit = payload[1] & 0x80;
                int endBit = payload[1] & 0x40;
                
                if (startBit) {
                    // 0xE0 ： 1110 0000
                    // 0x1F ： 0001 1111
                    // 重组NAL头
                    unsigned char nalHeader = (payload[0] & 0xE0) | (payload[1] & 0x1F);
                    outFile.write(startcode, sizeof(startcode));
                    outFile.write((char*)&nalHeader, 1);
                    outFile.write((char*)(payload + 2), recvLen - payloadOffset - 2);
                } else {
                    outFile.write((char*)(payload + 2), recvLen - payloadOffset - 2);
                }
            }
        }
    }


    // 清理资源
    outFile.close();
    closesocket(clientSockfd);
    closesocket(rtpSockfd);
    closesocket(rtcpSockfd);
    return 0;
}